from functools import partial

import numpy as np
import rclpy
from rclpy.node import Node
from rclpy.qos import QoSProfile
from std_msgs.msg import Float32

import matplotlib.pyplot as plt


items = [{"topic": "chatter", "range": 1, "offset": 0, "data": [], "line": None}]


def callback(msg, item):
    item["data"].append(msg.data * item["range"] + item["offset"])


def on_key(event):
    if event.key == "escape":  # 检测 ESC 按键
        rclpy.shutdown()


fig, ax = plt.subplots()
for item in items:
    (item["line"],) = ax.plot([], [], label=item["topic"])
fig.legend()
fig.canvas.mpl_connect("key_press_event", on_key)

rclpy.init()
node = Node("receiver")
for item in items:
    node.create_subscription(
        Float32, item["topic"], partial(callback, item=item), QoSProfile(depth=10)
    )


while rclpy.ok():
    rclpy.spin_once(node)
    for item in items:
        item["line"].set_xdata(np.arange(len(item["data"])))
        item["line"].set_ydata(item["data"])
    ax.relim()
    ax.autoscale_view()
    plt.draw()
    plt.pause(0.00001)
